# -*- coding: utf-8 -*-

"""
Mondrian 匿名化算法
@Time       :   2024/05/24
@Author     :   Li Kuangyuan
@Version    :   1.0
@Contact    :   1767958859@qq.com
@Software   :   VsCode
"""
import copy
import time

from functools import cmp_to_key
from util.data import column_change, column_restore
from util.types import NRange
from util.utility import *

class Partition(object):

    """用于分组, 保存记录
    self.member.: 记录
    self.width: 分区范围, 类别属性中是叶节点的数量, 数值属性中是数字范围
    self.middle: 泛化结果
    self.flag: 是否允许拆分
    """

    def __init__(self, data, range, result):
        """
        初始化数据
        """
        self.record = list(data)
        self.range = list(range)
        self.result = list(result)
        self.flag = [True] * QUASI_INDEX_LEN

    def __len__(self):
        return len(self.record)


def get_normalized_range(partition, index):
    # 返回分区的归一化范围
    if CATEGORY_FLAG_LIST[index]:   
        # 类别属性
        rangE = partition.range[index]
    else:   
        # 数值属性
        rangE = float(TAXONOMY_TREES[index].sort_value[partition.range[index][1]]) - float(TAXONOMY_TREES[index].sort_value[partition.range[index][0]])
    return rangE * 1.0 / QUASI_INDEX_RANGE[index]


def choose_dimension(partition):
    # 选择具有最大归一化范围的维度
    max_width = -1
    max_index = -1
    for i in range(QUASI_INDEX_LEN):
        if not partition.flag[i]:
            continue
        norm_range = get_normalized_range(partition, i)
        if norm_range > max_width:
            max_width = norm_range
            max_index = i
    return max_index

def find_mid(partition, dim):
    """
    找到分区的中间值
    """
    freq_dict = dict()
    for item in partition.record:
        freq_dict[item[dim]] = freq_dict.get(item[dim], 0) + 1

    total, values = sum(freq_dict.values()), list(freq_dict)
    middle = total / 2
    # 数量不足2*K或只有一种类别, 直接放回默认值
    if middle < MONDRAIN_K or len(values) <= 1:
        return '', '', values[0], values[-1]
    values.sort(key=cmp_to_key(cmp_num))
    
    count = 0
    split_index = 0
    split_val = ''
    # 逐个遍历找中位数对应的类别
    for i, key in enumerate(values):
        count += freq_dict[key]
        if count >= middle:
            split_val = key
            split_index = i
            break
    if len(values) > split_index + 1:
        next_val = values[split_index + 1]
    else:
        next_val = split_val
    return (split_val, next_val, values[0], values[-1])


def split_num_value(value_range, split_point_val):
    """
    以split_point_val拆分数值范围, 返回以split_point_val分割出的两个范围
        参数说明:
            range_value: 值的范围
            split_point_val: 以此值分割范围
    """
    rangE = value_range.split(',')
    if len(rangE) > 1:
        if rangE[0] == split_point_val:
            left = rangE[0]
        else:
            left = rangE[0] + ',' + split_point_val
        if rangE[1] == split_point_val:
            right = rangE[1]
        else:
            right = split_point_val + ',' + rangE[1]
        return left, right
    else:
       return rangE[0], rangE[0]


def split_numerical(partition, dim_index, par_range, par_result):
    # 通过中值来分割数值属性

    # 寻找分割点
    spilt_val, next_val, min_value, max_value = find_mid(partition, dim_index)
    # 更新范围
    if min_value != max_value:
        par_result[dim_index] = min_value + ',' + max_value
    else:
        par_result[dim_index] = max_value
    # 更新数据范围
    par_range[dim_index] = (TAXONOMY_TREES[dim_index].dict[min_value], TAXONOMY_TREES[dim_index].dict[max_value])
    # 未找到分割点, 无法分割
    if spilt_val == '' or spilt_val == next_val:
        return []
    # 以中值拆分对应分区的泛化结果
    left_result, right_result = par_result[:], par_result[:]
    left_result[dim_index], right_result[dim_index] = split_num_value(par_result[dim_index], spilt_val)
    
    # 按中间值的位置分割数据成两部分
    middle_val_pos = TAXONOMY_TREES[dim_index].dict[spilt_val]
    left_par_data, right_par_data = [], []
    for item in partition.record:
        pos = TAXONOMY_TREES[dim_index].dict[item[dim_index]]
        if pos > middle_val_pos:
            right_par_data.append(item)
        else:
            left_par_data.append(item)
    # 按中间值的位置更新左右分区范围
    left_range, right_range = par_range[:], par_range[:]
    left_range[dim_index], right_range[dim_index] = (par_range[dim_index][0], middle_val_pos), (TAXONOMY_TREES[dim_index].dict[next_val], par_range[dim_index][1])
   
    # 构建两个子分区
    res = []
    res.append(Partition(left_par_data, left_range, left_result))
    res.append(Partition(right_par_data, right_range, right_result))
    return res

def split_categorical(partition, dim_index, par_range, par_result):
    # 通过中值来分割类别属性
    # 分割
    split_val = TAXONOMY_TREES[dim_index][partition.result[dim_index]]
    sub_node = [t for t in split_val.child]
    if len(sub_node) == 0:
        return []
    sub_groups = [[] for _ in range(len(sub_node))]

    # 遍历分区记录，更新节点
    for temp in partition.record:
        for i, node in enumerate(sub_node):
            if temp[dim_index] in node.cover.keys():
                sub_groups[i].append(temp)
    for _, sub_group in enumerate(sub_groups):
        if len(sub_group) == 0:
            continue
        elif len(sub_group) < MONDRAIN_K:
            return []
    res_list = []
    for i, sub_group in enumerate(sub_groups):
        if len(sub_group) == 0:
            continue
        # 提取泛化结果
        res_tmp = par_result[:]
        res_tmp[dim_index] = sub_node[i].value
        # 提取范围
        rge_tmp = par_range[:]
        rge_tmp[dim_index] = len(sub_node[i])
        res_list.append(Partition(sub_group, rge_tmp, res_tmp))
    return res_list
    

def split_partition_to_list(partition, dim):
    # 分割分区
    par_range = partition.range
    par_result = partition.result
    if CATEGORY_FLAG_LIST[dim]:    
        # 类别属性划分
        return split_categorical(partition, dim, par_range, par_result)
    else:
        # 数值属性划分
        return split_numerical(partition, dim, par_range, par_result)

def anonymize(partition):
    # 递归地对组进行分区
    if check_splitable(partition) is False:
        RESULT.append(partition)
        return
    else:
        # 选择具有最大归一化范围的维度
        dim = choose_dimension(partition)
        # 分区
        par_list = split_partition_to_list(partition, dim)
        if len(par_list) != 0:
            # 对子分区继续划分
            for par in par_list:
                anonymize(par)
        else:
            partition.flag[dim] = 0
            # 递归分区其他维度
            anonymize(partition)

def check_splitable(partition):
    # 检查分区是否可以在满足k匿名性的情况下进一步拆分。
    return sum(partition.flag) != 0

def mondrian(att_trees, data, k, quasi_index_num, not_quasi_index):
    # Mondrian匿名化算法

    # 初始化所有需要使用的全局变量
    global MONDRAIN_K, RESULT, QUASI_INDEX_LEN, TAXONOMY_TREES, QUASI_INDEX_RANGE, CATEGORY_FLAG_LIST, NOT_QUASI_INDEX
    TAXONOMY_TREES = att_trees
    MONDRAIN_K = k
    RESULT = []
    CATEGORY_FLAG_LIST = []
    QUASI_INDEX_LEN, NOT_QUASI_INDEX, QUASI_INDEX_RANGE = quasi_index_num, not_quasi_index, [] # 准标识符和非准标识符的初始化
    for t in att_trees:
        CATEGORY_FLAG_LIST.append(not isinstance(t, NRange))

    
    # 统计每列的范围和中值, 构建分区
    middle_result = []
    range_temp = []
    for i in range(QUASI_INDEX_LEN):
        if CATEGORY_FLAG_LIST[i]:
            # 类别属性
            range_temp.append(len(TAXONOMY_TREES[i]['*']))
            QUASI_INDEX_RANGE.append(len(TAXONOMY_TREES[i]['*']))
            middle_result.append('*')
        else:
            # 数值属性         
            range_temp.append((0, len(TAXONOMY_TREES[i].sort_value) - 1))
            QUASI_INDEX_RANGE.append(TAXONOMY_TREES[i].range)
            middle_result.append(TAXONOMY_TREES[i].value)
    partitions = Partition(data, range_temp, middle_result)
    
    # 细分 分区
    anonymize(partitions)
    result = calculate_result()
    return result

def calculate_result():
    result = []
    for partition in RESULT:
        for i in range(len(partition)):
            # 添加非标识符属性
            not_quasi_list = []
            for s in range(len(partition.record[i]) - len(NOT_QUASI_INDEX), len(partition.record[i])):
                not_quasi_list = not_quasi_list + [partition.record[i][s]]
            not_quasi_list = partition.result + not_quasi_list
            result.append(not_quasi_list)
    return result



def mon_get_result_one(att_trees, data, k, quasi_index, not_quasi_index):
    result = mondrian(att_trees, 
            column_change(copy.deepcopy(data), quasi_index), k, len(quasi_index), not_quasi_index)
    return column_restore(result, quasi_index)